package com.alex.space.commons.tree;

/**
 * 平衡二叉树
 *
 * @author Alex
 * Created by Alex on 2018/3/13.
 */
public class AVLTree {

	/**
	 * 左高
	 */
	private static final int LH = 1;
	/**
	 * 等高
	 */
	private static final int EH = 0;
	/**
	 * 右高
	 */
	private static final int RH = -1;

	public AVLTreeNode root;

	public int size = 0;

	public void init() {
		root = null;
	}

	/**
	 * * p 最小旋转子树的根节点
	 * * 向左旋转之后p移到B的左子树处，p的右子树B变为此最小子树根节点，
	 * * B的左子树变为p的右子树
	 * * 比如：     A(-2)                   B(1)
	 * *             \        左旋转       /   \
	 * *            B(0)     ---->       A(0)   \
	 * *            /   \                   \    \
	 * *          BL(0)  BR(0)              BL(0) BR(0)
	 * * 旋转之后树的深度之差不超过1
	 */
	private void rotateLeft(AVLTreeNode p) {
		if (p != null) {
			System.out.println(p.data + "左旋");
			AVLTreeNode r = p.right;
			//把p右子树的左节点嫁接到p的右节点
			p.right = r.left;
			if (r.left != null) {
				//如果B的左节点BL不为空，把BL的父节点设为A
				r.left.parent = p;
			}
			//A的父节点设为B的父节点
			r.parent = p.parent;

			if (p.parent == null) {
				//如果p是根节点
				//r变为父节点，即B为父节点
				root = r;
			} else if (p.parent.left == p) {
				//如果p是左子节点
				//p的父节点的左子树为r
				p.parent.left = r;
			} else {
				//如果p是右子节点
				//p的父节点的右子树为r
				p.parent.right = r;
			}
			//p变为r的左子树，即A为B的左子树
			r.left = p;
			//同时更改p的父节点为r，即A的父节点为B
			p.parent = r;
		}
	}

	/**
	 * * p 最小旋转子树的根节点
	 * * 向右旋转之后，p移到B的右子节点处，p的左子树B变为最小旋转子树的根节点
	 * * B的右子节点变为p的左节点
	 * * 例如:   A(2)                     B(-1)
	 * *        /         右旋转          /    \
	 * *       B(0)       ------>       /     A(0)
	 * *      /   \                    /      /
	 * *   BL(0) BR(0)              BL(0)  BR(0)
	 */
	private void rotateRight(AVLTreeNode p) {
		if (p != null) {
			System.out.println(p.data + "右旋");
			AVLTreeNode l = p.left;
			//把B的右节点BR作为A的左节点
			p.left = l.right;
			if (l.right != null) {
				//如果BR不为null，设置BR的父节点为A
				l.right.parent = p;
			}
			//A的父节点赋给B的父节点
			l.parent = p.parent;
			if (p.parent == null) {
				//如果p是根节点
				//B为根节点
				root = l;
			} else if (p.parent.right == p) {
				//如果A是其父节点的左子节点
				//B为A的父节点的左子树
				p.parent.right = l;
			} else {
				//如果A是其父节点的右子节点
				//B为A的父节点的右子树
				p.parent.left = l;
			}
			//A为B的右子树
			l.right = p;
			//设置A的父节点为B
			p.parent = l;
		}
	}

	/**
	 * 插入操作
	 * 1.首先如同二叉排序树一般，找到其要插入的节点的位置，并把元素插入其中；
	 * 2.自下向上进行回溯，回溯做两个事情：
	 * (1)修改祖先节点的平衡因子，当插入一个节点时根节点到插入节点的路径中的节点的平衡因子会被改变
	 * (2)在改变祖先节点的平衡因子的同时，找到最近一个平衡因子大于2或小于-2的节点，从这个节点开始调整最小不平衡树进行旋转调整
	 * 3.有几种情况
	 * a. 树左边高，右边低，平衡因子符号相同，则右旋。
	 * b. 树右边高，左边低，平衡因子符号相同，则左旋。
	 * c. 树右边高，左边低，平衡因子符号不同，先右旋再左旋（转变成b）
	 * d. 和d相反，左边高，右边低，平衡因子符号不同，先左旋再右旋（转变成a）。
	 */
	public boolean add(int elem) {
		AVLTreeNode t = root;
		if (t == null) {
			root = new AVLTreeNode();
			root.data = elem;
			root.parent = null;
			size = 1;
			return true;
		}

		AVLTreeNode parent;  //保存t的父节点
		boolean isLeftChild = true;
		//从根节点向下搜索，找到插入位置
		do {
			parent = t;

			if (elem < t.data) {
				t = t.left;
				isLeftChild = true;
			} else if (elem > t.data) {
				t = t.right;
				isLeftChild = false;
			} else {
				return false;
			}
		} while (t != null);

		AVLTreeNode child = new AVLTreeNode();
		child.data = elem;
		child.parent = parent;
		if (isLeftChild) {
			parent.left = child;

		} else {
			parent.right = child;
		}

		//自下向上回溯，查找最近不平衡节点
		while (parent != null) {
			//插入节点在parent的左子树中
			if (elem < parent.data) {
				parent.bf++;
			} else {
				//插入节点在parent的右子树中
				parent.bf--;
			}
			if (parent.bf == 0) {
				//此节点的balance为0，不再向上调整BF值，且不需要旋转
				break;
			}

			//找到最小不平衡子树根节点
			if (Math.abs(parent.bf) == 2) {
				fixAfterInsertion(parent);
				break;                  //不用继续向上回溯
			}
			parent = parent.parent;
		}
		size++;
		return true;
	}

	/**
	 * 调整的方法：
	 * 1.当树的根bf为2时，即左子树高于右子树：
	 * 如果R的左子树的根节点的BF为1时，做右旋；
	 * 如果R的左子树的根节点的BF为-1时，先左旋然后再右旋
	 * <p>
	 * 2.当树的根bf为-2时，即右子树高于左子树：
	 * 如果R的右子树的根节点的BF为1时，先右旋后左旋
	 * 如果R的右子树的根节点的BF为-1时，做左旋
	 */
	private void fixAfterInsertion(AVLTreeNode p) {
		if (p.bf == 2) {
			leftBalance(p);
		}
		if (p.bf == -2) {
			rightBalance(p);
		}
	}

	/**
	 * * 做左平衡处理
	 * * 平衡因子的调整如图：
	 * *         t                       rd
	 * *       /   \                   /    \
	 * *      l    tr   左旋后右旋    l       t
	 * *    /   \       ------->    /  \       \
	 * *  ll    rd                ll   rdl     tr
	 * *       /
	 * *     rdl
	 * *
	 * *   情况1(rd的BF为1)
	 * *
	 * *        t                       rd
	 * *      /   \                   /    \
	 * *     l    tr   左旋后右旋    l       t
	 * *   /   \       ------->    /  \    /  \
	 * * ll    rd                ll   rdl rdr  tr
	 * *      /   \
	 * *    rdl  rdr
	 * *
	 * *   情况2(rd的BF为0)
	 * *
	 * *
	 * *         t                       rd
	 * *       /   \                   /    \
	 * *      l    tr   左旋后右旋    l       t
	 * *    /   \       ------->    /       /  \
	 * *  ll    rd                ll       rdr  tr
	 * *           \
	 * *          rdr
	 * *
	 * *   情况3(rd的BF为-1)
	 * *
	 * *
	 * *         t                         l
	 * *       /       右旋处理          /    \
	 * *      l        ------>          ll     t
	 * *    /   \                             /
	 * *   ll   rd                           rd
	 * *   情况4(L等高)
	 */
	private boolean leftBalance(AVLTreeNode t) {
		boolean heightLower = true;
		AVLTreeNode l = t.left;
		switch (l.bf) {
			case LH:
				//一般情况，左高，右旋调整，旋转后树的高度减小
				t.bf = l.bf = EH;
				rotateRight(t);
				break;
			case RH:
				//右高，分情况调整
				AVLTreeNode rd = l.right;
				switch (rd.bf) {
					//调整各个节点的BF
					case LH:
						//情况1
						t.bf = RH;
						l.bf = EH;
						break;
					case EH:    //情况2
						t.bf = l.bf = EH;
						break;
					case RH:    //情况3
						t.bf = EH;
						l.bf = LH;
						break;
					default:
						break;
				}
				rd.bf = EH;
				rotateLeft(t.left);
				rotateRight(t);
				break;
			case EH:
				//特殊情况4，这种情况在添加时不可能出现，只在移除时可能出现，旋转之后整体树高不变
				l.bf = RH;
				t.bf = LH;
				rotateRight(t);
				heightLower = false;
				break;
			default:
				break;
		}
		return heightLower;
	}

	/**
	 * * 做右平衡处理
	 * * 平衡因子的调整如图：
	 * *           t                               ld
	 * *        /     \                          /     \
	 * *      tl       r       先右旋再左旋     t       r
	 * *             /   \     -------->      /   \    /  \
	 * *           ld    rr                 tl   ldl  ldr rr
	 * *          /  \
	 * *       ldl  ldr
	 * *       情况2(ld的BF为0)
	 * *
	 * *           t                               ld
	 * *        /     \                          /     \
	 * *      tl       r       先右旋再左旋     t       r
	 * *             /   \     -------->      /   \       \
	 * *           ld    rr                 tl   ldl      rr
	 * *          /
	 * *       ldl
	 * *       情况1(ld的BF为1)
	 * *
	 * *           t                               ld
	 * *        /     \                          /     \
	 * *      tl       r       先右旋再左旋     t       r
	 * *             /   \     -------->      /        /  \
	 * *           ld    rr                 tl        ldr rr
	 * *             \
	 * *             ldr
	 * *       情况3(ld的BF为-1)
	 * *
	 * *           t                                  r
	 * *             \          左旋                /   \
	 * *              r        ------->           t      rr
	 * *            /   \                          \
	 * *           ld   rr                         ld
	 * *        情况4(r的BF为0)
	 */
	private boolean rightBalance(AVLTreeNode t) {
		boolean heightLower = true;
		AVLTreeNode r = t.right;
		switch (r.bf) {
			case LH:
				//左高，分情况调整
				AVLTreeNode ld = r.left;
				switch (ld.bf) {
					//调整各个节点的BF
					case LH:
						//情况1
						t.bf = EH;
						r.bf = RH;
						break;
					case EH:
						//情况2
						t.bf = r.bf = EH;
						break;
					case RH:
						//情况3
						t.bf = LH;
						r.bf = EH;
						break;
					default:
						break;
				}
				ld.bf = EH;
				rotateRight(t.right);
				rotateLeft(t);
				break;
			case RH:
				//右高，左旋调整
				t.bf = r.bf = EH;
				rotateLeft(t);
				break;
			case EH:
				//特殊情况4
				r.bf = LH;
				t.bf = RH;
				rotateLeft(t);
				heightLower = false;
				break;
			default:
				break;
		}
		return heightLower;
	}

	/**
	 * 查找指定元素，如果找到返回其对象，否则返回null
	 */
	private AVLTreeNode getEntry(int elem) {
		AVLTreeNode tmp = root;
		int c;
		while (tmp != null) {
			if (elem == tmp.data) {
				return tmp;
			} else if (elem < tmp.data) {
				tmp = tmp.left;
			} else {
				tmp = tmp.right;
			}
		}
		return null;
	}

	/**
	 * 平衡二叉树的移除元素操作
	 */
	public boolean remove(int elem) {
		AVLTreeNode e = getEntry(elem);
		if (e != null) {
			deleteEntry(e);
			return true;
		}
		return false;
	}

	private void deleteEntry(AVLTreeNode p) {
		size--;
		//如果p左右子树都不为空，找到其直接后继，替换p，之后p指向s，删除p其实是删除s
		//所有的删除左右子树不为空的节点都可以调整为删除左右子树有其一不为空，或都为空的情况。
		if (p.left != null && p.right != null) {
			AVLTreeNode s = successor(p);
			p.data = s.data;
			p = s;
		}
		AVLTreeNode replacement = (p.left != null ? p.left : p.right);

		//如果其左右子树有其一不为空
		if (replacement != null) {
			replacement.parent = p.parent;
			if (p.parent == null) {
				//如果p为root节点
				root = replacement;
			} else if (p == p.parent.left) {
				//如果p是左孩子
				p.parent.left = replacement;
			} else {
				//如果p是右孩子
				p.parent.right = replacement;
			}

			//p的指针清空
			p.left = p.right = p.parent = null;

			//这里更改了replacement的父节点，所以可以直接从它开始向上回溯
			fixAfterDeletion(replacement);

		}

		// 如果全树只有一个节点
		else if (p.parent == null) {
			root = null;
		}

		//左右子树都为空
		else {
			//这里从p开始回溯
			fixAfterDeletion(p);
			if (p.parent != null) {
				if (p == p.parent.left) {
					p.parent.left = null;
				} else if (p == p.parent.right) {
					p.parent.right = null;
				}

				p.parent = null;
			}
		}

	}

	/**
	 * 返回以中序遍历方式遍历树时，t的直接后继
	 */
	private AVLTreeNode successor(AVLTreeNode t) {
		if (t == null) {
			return null;
		}

		//往右，然后向左直到尽头
		else if (t.right != null) {
			AVLTreeNode p = t.right;
			while (p.left != null) {
				p = p.left;
			}
			return p;
		}

		//right为空，如果t是p的左子树，则p为t的直接后继
		else {
			AVLTreeNode p = t.parent;
			AVLTreeNode ch = t;
			while (p != null && ch == p.right) {
				;//如果t是p的右子树，则继续向上搜索其直接后继
				ch = p;
				p = p.parent;
			}
			return p;
		}

	}

	/**
	 * 删除某节点p后的调整方法
	 * 1.从p开始向上回溯，修改祖先的BF值，这里只要调整从p的父节点到根节点的BF值
	 * 调整原则为，当p位于某祖先节点(简称A)的左子树中时，A的BF减1
	 * 当p位于A的，右子树中时A的BF加1。当某个祖先节点BF变为1或-1时停止回溯，这里与插入是相反的
	 * 因为原本这个节点是平衡的，删除它的子树的某个节点并不会改变它的高度
	 * 2.检查每个节点的BF值，如果为2或-2需要进行旋转调整
	 */
	private void fixAfterDeletion(AVLTreeNode p) {
		//看最小子树调整后，它的高度是否发生变化，如果减小，继续回溯
		boolean heightLower = true;
		AVLTreeNode t = p.parent;
		//自下向上回溯，查找不平衡的节点进行调整
		while (t != null && heightLower) {
			/*
			 * 删除的节点是右子树，等于的话，必然是删除的某个节点的左右子树不为空的情况
			 * 例如：     10
			 *          /    \
			 *         5     15
			 *       /   \
			 *      3    6
			 * 这里删除5，是把6的值赋给5，然后删除6，这里6是p，p的父节点的值也是6。
			 * 而这也是右子树的一种
			 */
			if (p.data >= t.data) {
				t.bf++;
			} else {
				t.bf--;
			}

			//父节点经过调整平衡因子后，如果为1或-1，说明调整之前是0，停止回溯。
			if (Math.abs(t.bf) == 1) {
				break;
			}
			AVLTreeNode r = t;
			//这里的调整跟插入一样
			if (t.bf == 2) {
				heightLower = leftBalance(r);
			} else if (t.bf == -2) {
				heightLower = rightBalance(r);
			}
			t = t.parent;
		}

	}

}
